package utils

import (
	"bytes"
	"crypto/cipher"
	"crypto/md5"
	"encoding/hex"
	"github.com/tjfoc/gmsm/sm4"
)

type SM4EncryptStruct struct {
	SM4Key string
	SM4Iv  string
}

func GetSM4EncryptInstance(SM4Key string, SM4Iv string) *SM4EncryptStruct {
	return &SM4EncryptStruct{
		SM4Key: SM4Key,
		SM4Iv:  SM4Iv,
	}
}

func (SM4 *SM4EncryptStruct) Encrypt(data string) (result string, err error) {
	plainText := []byte(data)
	iv := []byte(SM4.SM4Iv)
	key := []byte(SM4.SM4Key)
	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}
	paddingData := paddingLastGroup(plainText, block.BlockSize())
	blockMode := cipher.NewCBCEncrypter(block, iv)
	cipherText := make([]byte, len(paddingData))
	blockMode.CryptBlocks(cipherText, paddingData)
	cipherString := hex.EncodeToString(cipherText)
	return cipherString, nil
}

func (SM4 *SM4EncryptStruct) Decrypt(data string) (res string, err error) {

	iv := []byte(SM4.SM4Iv)
	key := []byte(SM4.SM4Key)
	block, err := sm4.NewCipher(key)
	if err != nil {
		return "", err
	}

	decodeString, err := hex.DecodeString(data)
	if err != nil {
		return "", err
	}
	blockMode := cipher.NewCBCDecrypter(block, iv)
	blockMode.CryptBlocks(decodeString, decodeString)
	plainText := unPaddingLastGroup(decodeString)
	return string(plainText), nil
}

func paddingLastGroup(plainText []byte, blockSize int) []byte {
	padNum := blockSize - len(plainText)%blockSize
	char := []byte{byte(padNum)}
	newPlain := bytes.Repeat(char, padNum)
	newText := append(plainText, newPlain...)
	return newText
}

func unPaddingLastGroup(plainText []byte) []byte {
	length := len(plainText)
	lastChar := plainText[length-1]
	number := int(lastChar)
	return plainText[:length-number]
}

func Md5(data []byte) string {
	hash := md5.New()
	hash.Write([]byte(data))
	hashBytes := hash.Sum(nil)
	return hex.EncodeToString(hashBytes)
}
